package com.utils;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.spec.KeySpec;
import java.util.Base64;

/**
 * @BelongsProject: study
 * @BelongsPackage: com.utils
 * @Author: zhangzhanheng
 * @CreateTime: 2024-01-17  20:36
 * @Description: TODO
 * @Version: 1.0
 */
public class StringEncryptorDecryptor {
    private static final String ENCRYPTION_ALGORITHM = "AES";
    private static final String SECRET_KEY_ALGORITHM = "PBKDF2WithHmacSHA256";
    private static final int ITERATION_COUNT = 65536;
    private static final int KEY_LENGTH = 256;
    private static final String CIPHER_TRANSFORMATION = "AES/CBC/PKCS5Padding";

    public static String encrypt(String input, String key) throws Exception {
        SecretKey secretKey = generateKey(key);
        Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);

        // Generate IV
        byte[] ivBytes = new byte[cipher.getBlockSize()];
        IvParameterSpec ivSpec = new IvParameterSpec(ivBytes);

        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivSpec);

        byte[] paddedInput = addPadding(input.getBytes(StandardCharsets.UTF_8));
        byte[] encryptedBytes = cipher.doFinal(paddedInput);
        String base64Encoded = Base64.getEncoder().encodeToString(encryptedBytes);

        // Remove non-alphanumeric characters
        return base64Encoded.replaceAll("[^a-zA-Z0-9]", "");
    }

    public static String decrypt(String encryptedInput, String key) throws Exception {
        SecretKey secretKey = generateKey(key);
        Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);

        // Extract IV from the encrypted string
        byte[] ivBytes = new byte[cipher.getBlockSize()];
        IvParameterSpec ivSpec = new IvParameterSpec(ivBytes);

        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivSpec);

        // Base64 decode and then decrypt
        byte[] encryptedBytes = Base64.getDecoder().decode(encryptedInput);
        byte[] decryptedBytes = cipher.doFinal(encryptedBytes);

        return new String(removePadding(decryptedBytes), StandardCharsets.UTF_8);
    }

    private static SecretKey generateKey(String key) throws Exception {
        SecretKeyFactory factory = SecretKeyFactory.getInstance(SECRET_KEY_ALGORITHM);
        KeySpec spec = new PBEKeySpec(key.toCharArray(), key.getBytes(), ITERATION_COUNT, KEY_LENGTH);
        SecretKey tmp = factory.generateSecret(spec);
        return new SecretKeySpec(tmp.getEncoded(), ENCRYPTION_ALGORITHM);
    }

    private static byte[] addPadding(byte[] input) {
        int paddingLength = 16 - input.length % 16;
        byte[] paddedInput = new byte[input.length + paddingLength];
        System.arraycopy(input, 0, paddedInput, 0, input.length);
        return paddedInput;
    }

    private static byte[] removePadding(byte[] input) {
        int paddingLength = 0;
        for (int i = input.length - 1; i >= 0; i--) {
            if (input[i] == 0) {
                paddingLength++;
            } else {
                break;
            }
        }
        byte[] unpaddedInput = new byte[input.length - paddingLength];
        System.arraycopy(input, 0, unpaddedInput, 0, unpaddedInput.length);
        return unpaddedInput;
    }

    public static void main(String[] args) {
        try {
            String originalString = "httcs0702101";
            String key = "yourSecretKey"; // 你的密钥

            // 加密
            String encryptedString = encrypt(originalString, key);
            System.out.println("Encrypted: " + encryptedString);

            // 解密
            String decryptedString = decrypt(encryptedString, key);
            System.out.println("Decrypted: " + decryptedString);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
